package world

import (
	"MetroBuilderGameJam/gameEngine/utils"
	"fmt"
	"image/color"
	"reflect"
	"slices"
)

func NewRailRoad(city *City, color color.Color, available bool) *RailRoad {
	return &RailRoad{
		branches:    make([]*BranchCompleted, 0),
		color:       color,
		isAvailable: available,
		city:        city,
	}
}

type RailRoad struct {
	city               *City
	color              color.Color
	branches           []*BranchCompleted
	loopConnection     bool
	isAvailable        bool
	_stations          map[*Station]struct{}
	_cachedConnections []*RailRoadConnection
}

func (rr *RailRoad) IsAvailable() bool {
	return rr.isAvailable
}

func (rr *RailRoad) AddBranch(b *BranchCompleted) error {
	if b == nil {
		return nil
	}

	if rr.loopConnection {
		return ErrRoadIsLooped{}
	}

	err := rr.checkBranchIsNew(b)
	if err != nil {
		return err
	}
	if rr.GetLastEdge() == nil || b.from == rr.GetLastEdge().station {
		rr.branches = append(rr.branches, b)
	} else if b.from == rr.GetFirstEdge().station {
		rr.branches = append([]*BranchCompleted{b.reverse()}, rr.branches...)
	} else {
		return ErrBranchNotConnectedToAnyEdges{}
	}
	b.AttachToRoad(rr)
	fe := rr.GetFirstEdge()
	le := rr.GetLastEdge()
	rr.loopConnection = rr.Count() > 0 && fe != nil && le != nil && fe.station == le.station
	rr.recalcStations()

	return nil
}

func (rr *RailRoad) ExcludeBranch(branch *BranchCompleted) error {
	for i, b := range rr.branches {
		if b != branch {
			continue
		}
		rr.branches = append(rr.branches[:i], rr.branches[i+1:]...)
		fe := rr.GetFirstEdge()
		le := rr.GetLastEdge()
		rr.loopConnection = rr.Count() > 0 && fe != nil && le != nil && fe.station == le.station
		rr.recalcStations()
		return nil
	}

	return fmt.Errorf("branch not found")
}

func (rr *RailRoad) GetColor() color.Color {
	return rr.color
}

func (rr *RailRoad) IsLoop() bool {
	return rr.loopConnection
}

func (rr *RailRoad) checkBranchIsNew(b *BranchCompleted) error {
	for _, branch := range rr.branches {
		if branch.from == b.from && branch.to == b.to || branch.from == b.to && branch.to == b.from {
			return ErrBranchAlreadyExist{}
		}
	}
	return nil
}

func (rr *RailRoad) Count() int {
	return len(rr.branches)
}

func (rr *RailRoad) DisconnectStation(station *Station) (nextStation *Station, err error) {
	for _, e := range []*BranchEdge{rr.GetFirstEdge(), rr.GetLastEdge()} {
		if e.station == station {
			err = rr.ExcludeBranch(e.branch)
			if err != nil {
				return nil, err
			}
			if e.branch.from == station {
				return e.branch.to, nil
			}
			return e.branch.from, nil
		}
	}
	for i := 1; i < len(rr.branches); i++ {
		b := rr.branches[i]
		if b.from == station {
			rr.branches[i-1].to = b.to
			return nil, rr.ExcludeBranch(b)
		}
		if b.to == station && i < len(rr.branches)-1 {
			b.to = rr.branches[i+1].to
			return nil, rr.ExcludeBranch(rr.branches[i+1])
		}
	}
	return nil, fmt.Errorf("station not found")
}

func (rr *RailRoad) findBranchIndexForStation(station *Station, startIndex int) (int, error) {
	for i := startIndex; i < len(rr.branches); i++ {
		if rr.branches[i].from == station || rr.branches[i].to == station {
			return i, nil
		}
	}
	return 0, fmt.Errorf("branch for station not found")
}

func (rr *RailRoad) GetFirstEdge() *BranchEdge {
	if len(rr.branches) == 0 {
		return nil
	}
	return NewBranchEdge(rr.branches[0].from, rr.branches[0])

}

func (rr *RailRoad) GetLastEdge() *BranchEdge {
	if len(rr.branches) == 0 {
		return nil
	}
	lastBranch := rr.branches[len(rr.branches)-1]
	return NewBranchEdge(lastBranch.to, lastBranch)
}

func (rr *RailRoad) AddStationToCompletedBranch(branch *BranchCompleted, station *Station) error {
	if branch == nil || station == nil {
		return nil
	}
	if rr.HasStation(station) {
		return fmt.Errorf("station already exist")
	}
	for i, b := range rr.branches {
		if b == branch {
			bfrom, err := NewCompletedBranch(b.from, station, branch.road)
			if err != nil {
				return err
			}
			bto, err := NewCompletedBranch(station, b.to, branch.road)
			if err != nil {
				return err
			}
			other := append([]*BranchCompleted{bfrom, bto}, rr.branches[i+1:]...)
			bfrom.AttachToRoad(rr)
			bto.AttachToRoad(rr)
			rr.branches = append(rr.branches[:i], other...)
			rr.loopConnection = rr.Count() > 0 && rr.GetFirstEdge().station == rr.GetLastEdge().station
			return nil
		}
	}

	return fmt.Errorf("branch not found")
}

func (rr *RailRoad) HasStation(station *Station) bool {
	_, ok := rr._stations[station]
	return ok
}

func (rr *RailRoad) GetBranchForStation(station *Station) *BranchCompleted {
	for _, b := range rr.branches {
		if b.from == station || b.to == station {
			return b
		}
	}

	return nil
}

func (rr *RailRoad) GetAnotherCompletedBranchForStation(station *Station, branch *BranchCompleted) *BranchCompleted {
	for _, b := range rr.branches {
		if reflect.DeepEqual(b, branch) {
			continue
		}
		if b.from == station || b.to == station {
			return b
		}
	}

	return nil
}

func (rr *RailRoad) IsEdgedStation(station *Station) bool {
	return rr.GetFirstEdge().station == station || rr.GetLastEdge().station == station
}

func (rr *RailRoad) ExcludeExistsStationByProject(station *Station, branchProject *BranchProject) {
	nextStation, err := rr.DisconnectStation(station)
	if err == nil && nextStation != nil {
		branchProject.SetFromStation(nextStation)
	}
}

func (rr *RailRoad) HasBranch(branch RailRoadBranch) bool {
	for _, b := range rr.branches {
		if b.IsSameWith(branch.(*BranchCompleted)) {
			return true
		}
	}
	return false
}

func (rr *RailRoad) BuildPathBetweenStations(sFrom, sTo *Station) []*Station {
	if !rr.HasStation(sTo) || !rr.HasStation(sFrom) {
		return []*Station{}
	}
	br := rr.GetBranchForStation(sFrom)
	path := rr.BuildPathBetweenStationsFromBranch(sFrom, sTo, br)
	if len(path) == 0 {
		path = rr.BuildPathBetweenStationsFromBranch(sFrom, sTo, rr.GetAnotherCompletedBranchForStation(sFrom, br))
	}
	return path
}

func (rr *RailRoad) BuildPathBetweenStationsFromBranch(sFrom, sGoal *Station, brFrom *BranchCompleted) []*Station {
	if !rr.HasStation(sGoal) || !rr.HasStation(sFrom) || brFrom.RailRoad().GetColor() != rr.GetColor() || rr.Count() == 0 {
		return []*Station{}
	}
	path := make([]*Station, 0, rr.Count())

	var stTo *Station
	var brTo *BranchCompleted
	brFrom2 := rr.GetAnotherCompletedBranchForStation(sFrom, brFrom)
	if brFrom2 != nil {
		brFrom = brFrom2
	}
	bStart := *brFrom

	for stTo != sGoal {
		brTo = rr.GetAnotherCompletedBranchForStation(sFrom, brFrom)
		if brTo == nil || bStart.IsSameWith(brTo) {
			s := bStart.GetAnotherStation(sFrom)
			fe := rr.GetFirstEdge()
			le := rr.GetLastEdge()
			if s == sGoal || fe != nil && fe.FromStation() == s || le != nil && le.FromStation() == s {
				return []*Station{s}
			}
			break
		}
		stTo = brTo.GetAnotherStation(sFrom)
		path = append(path, stTo)
		sFrom = stTo
		brFrom = brTo
	}

	if len(path) == 0 || path[len(path)-1] != sGoal {
		return []*Station{}
	}

	return path
}

func (rr *RailRoad) BuildPathToStationShape(sFrom *Station, shape StationType) []*Station {
	if !rr.HasStation(sFrom) {
		return []*Station{}
	}
	br := rr.GetBranchForStation(sFrom)
	path := rr.BuildPathToStationShapeFromBranch(sFrom, shape, br)
	path2 := rr.BuildPathToStationShapeFromBranch(sFrom, shape, rr.GetAnotherCompletedBranchForStation(sFrom, br))
	if len(path2) == 0 {
		return path
	}
	if len(path) == 0 {
		return path2
	}
	if len(path2) < len(path) {
		return path2
	}
	return path
}

func (rr *RailRoad) BuildPathToStationShapeFromBranch(sFrom *Station, shape StationType, brFrom *BranchCompleted) []*Station {
	if !rr.HasStation(sFrom) || brFrom.RailRoad().GetColor() != rr.GetColor() || rr.Count() == 0 {
		return []*Station{}
	}
	path := make([]*Station, 0, rr.Count())

	var stTo *Station
	var brTo *BranchCompleted
	brFrom2 := rr.GetAnotherCompletedBranchForStation(sFrom, brFrom)
	if brFrom2 != nil {
		brFrom = brFrom2
	}
	bStart := *brFrom

	for stTo == nil || stTo.Shape() != shape {
		brTo = rr.GetAnotherCompletedBranchForStation(sFrom, brFrom)
		if brTo == nil || bStart.IsSameWith(brTo) {
			s := bStart.GetAnotherStation(sFrom)
			fe := rr.GetFirstEdge()
			le := rr.GetLastEdge()
			if s.Shape() == shape || fe != nil && fe.FromStation() == s || le != nil && le.FromStation() == s {
				return []*Station{s}
			}
			break
		}
		stTo = brTo.GetAnotherStation(sFrom)
		path = append(path, stTo)
		sFrom = stTo
		brFrom = brTo
	}

	if len(path) == 0 || path[len(path)-1].Shape() != shape {
		return []*Station{}
	}

	return path
}

func (rr *RailRoad) FindStationWithShape(shape StationType) *Station {
	for s := range rr._stations {
		if s.Shape() == shape {
			return s
		}
	}
	return nil
}

func (rr *RailRoad) ConnectionsWithStationShape(shape StationType) []*RailRoadConnection {
	return rr.connectionsWithStationShape(shape, []*RailRoad{})
}

func (rr *RailRoad) connectionsWithStationShape(shape StationType, excludeRoads []*RailRoad) []*RailRoadConnection {
	if rr._cachedConnections == nil {
		rr._cachedConnections = rr.city.findAllConnectionsDeep(rr, []*RailRoad{})
	}

	roads := make([]*RailRoad, 0)
	for _, c := range rr._cachedConnections {
		for _, r := range c.Roads {
			if r != rr && !slices.Contains(excludeRoads, r) {
				roads = append(roads, r)
			}
		}
	}

	var result []*RailRoadConnection
	for _, r := range roads {
		if r.FindStationWithShape(shape) != nil {
			for _, c := range rr._cachedConnections {
				if c.hasRoad(r) && c.hasRoad(rr) {
					result = append(result, c)
				}
			}
			continue
		}
		connections := r.connectionsWithStationShape(shape, append(excludeRoads, rr))
		if len(connections) > 0 {
			for _, c := range rr._cachedConnections {
				if c.hasRoad(r) && c.hasRoad(rr) {
					result = append(result, c)
				}
			}
		}
	}
	return result
}

func (rr *RailRoad) Stations() []*Station {
	return utils.MapKeys[*Station, struct{}](rr._stations)
}

func (rr *RailRoad) recalcStations() {
	rr._stations = make(map[*Station]struct{}, len(rr.branches)*2)

	for _, branch := range rr.branches {
		rr._stations[branch.from] = struct{}{}
		rr._stations[branch.to] = struct{}{}
	}
}

func (rr *RailRoad) forgotConnections() {
	rr._cachedConnections = nil
}

func NewRailRoadConnection(s *Station, road1, road2 *RailRoad) *RailRoadConnection {
	return &RailRoadConnection{
		Station: s,
		Roads:   []*RailRoad{road1, road2},
	}

}

type RailRoadConnection struct {
	Station *Station
	Roads   []*RailRoad
}

func (c RailRoadConnection) hasRoad(r *RailRoad) bool {
	for _, road := range c.Roads {
		if road == r {
			return true
		}
	}
	return false
}
